import pandas as pd
import joblib
import lightgbm as lgb
import numpy as np
from sklearn.model_selection import train_test_split


def create_features(df):
    """创建时间序列特征"""
    # # 滞后特征
    # df['lag1'] = df.groupby('spu_code')['sale_qty'].shift(1)
    # df['lag7'] = df.groupby('spu_code')['sale_qty'].shift(7)
    #
    # # 移动平均
    # df['rolling7_mean'] = df.groupby('spu_code')['sale_qty'].transform(
    #     lambda x: x.rolling(7, min_periods=1).mean().shift(1))
    #
    # df['rolling14_mean'] = df.groupby('spu_code')['sale_qty'].transform(
    #     lambda x: x.rolling(14, min_periods=1).mean().shift(1))

    # 时间衰减特征
    df['exp_decay'] = np.exp(-0.05 * df['days_since_launch'])

    # 周期性特征（每周第几天）
    # df['day_of_week'] = df['days_since_launch'] % 7

    return df.dropna()


def train_lightgbm_models(data_file):
    """使用LightGBM训练所有款号的模型并保存"""
    # 加载数据
    df = pd.read_excel(data_file)

    # 存储所有训练好的模型
    models = {}

    # 按款号分组训练
    for spu, group in df.groupby('spu_code'):
        # 创建特征
        group = create_features(group)

        # 准备训练数据
        features = ['sum_cart_total', 'new_sale_qty', 'days_since_launch',
                    # 'lag1', 'lag7', 'rolling7_mean', 'rolling14_mean',
                    'exp_decay'
            # , 'day_of_week'
                    ]

        X = group[features]
        y = group['sale_qty']

        # 划分训练集和验证集
        if len(X) > 20:
            X_train, X_val, y_train, y_val = train_test_split(
                X, y, test_size=0.2, random_state=42
            )
        else:
            # 数据量少时使用全部训练
            X_train, y_train = X, y
            X_val, y_val = X, y

        # 创建LightGBM数据集
        train_data = lgb.Dataset(X_train, label=y_train)
        val_data = lgb.Dataset(X_val, label=y_val, reference=train_data)

        # 设置模型参数 - 优化时间序列预测
        params = {
            'objective': 'regression',
            'metric': 'rmse',
            'learning_rate': 0.05,
            'num_leaves': 63,  # 增加复杂度
            'max_depth': 7,  # 增加深度
            'min_data_in_leaf': 5,
            'feature_fraction': 0.8,
            'bagging_fraction': 0.8,
            'bagging_freq': 5,
            'verbosity': -1,
            'seed': 42,
            'linear_tree': True,  # 更好处理趋势
            'extra_trees': True  # 减少过拟合
        }

        # 训练模型
        model = lgb.train(
            params,
            train_data,
            num_boost_round=1000,
            valid_sets=[val_data],
            callbacks=[
                lgb.early_stopping(stopping_rounds=50, verbose=False),
                lgb.log_evaluation(period=100)
            ]
        )

        # 存储模型
        models[spu] = (model, features)  # 同时保存特征列表
        print(f"款号 {spu} LightGBM模型训练完成，使用数据量: {len(group)}")

    # 保存所有模型
    joblib.dump(models, 'enhanced_lightgbm_models.pkl')
    print("所有模型已保存为 enhanced_lightgbm_models.pkl")

    return models


if __name__ == "__main__":
    models = train_lightgbm_models("enhanced_sales_data.xlsx")